﻿#include <stdio.h>
#include <stdint.h>
#include <pcap.h>
#include <string.h>
#include <threads.h>

#include "opts.h"
#include "bindings.h"
#include "catch.h"

#define SEND_DNS_PACKET_FAILED 1 // send DNS packet

uint16_t PACKET_COUNTER = 0;

int interactive_mode = false;

uint16_t S_PORT = 1234;


int start_capture(void *arg)
{
	cnd_t *cnd = (cnd_t *)arg;

	if (start(cnd, interactive_mode, S_PORT) != 0)
	{
		printf("start failed\n");
		thrd_exit(EXIT_FAILURE);
		return EXIT_FAILURE;
	}
	return 0;
}

/*
 * main 函数实现读取命令行参数。
 */
int main(int argc, char **argv)
{
	Config config = parse_config(argc, argv);

	printf("src_mac: %02X-%02X-%02X-%02X-%02X-%02X\n", config.src_mac[0], config.src_mac[1], config.src_mac[2], config.src_mac[3], config.src_mac[4], config.src_mac[5]);
	printf("src_ip: %d.%d.%d.%d\n", config.src_ip[0], config.src_ip[1], config.src_ip[2], config.src_ip[3]);
	printf("gtw_mac: %02X-%02X-%02X-%02X-%02X-%02X\n", config.gtw_mac[0], config.gtw_mac[1], config.gtw_mac[2], config.gtw_mac[3], config.gtw_mac[4], config.gtw_mac[5]);
	printf("dns_server: %d.%d.%d.%d\n", config.dns_server[0], config.dns_server[1], config.dns_server[2], config.dns_server[3]);
	printf("src_port: %d\n", config.s_port);

	S_PORT = config.s_port;

	mtx_t mtx;
	if (thrd_success != mtx_init(&mtx, mtx_plain))
	{
		printf("ERORR; mtx_init() call failed\n");
		exit(EXIT_FAILURE);
	}

	cnd_t cnd;
	if (thrd_success != cnd_init(&cnd))
	{
		printf("ERORR; cnd_init() call failed\n");
		exit(EXIT_FAILURE);
	}

	// 如果用户输入了域名，则直接发送 DNS 查询数据包，然后退出。
	if (config.domain_name != NULL)
	{
		// 发送 DNS 查询数据包

		uint16_t pkt_id = PACKET_COUNTER;
		PACKET_COUNTER++;

		Addresses addrs;

		// 没有处理 memcpy_s 可能出现的错误。
		memcpy_s(addrs.source_mac_addr, sizeof(addrs.source_mac_addr), config.src_mac, sizeof(config.src_mac));
		memcpy_s(addrs.source_ip_addr, sizeof(addrs.source_ip_addr), config.src_ip, sizeof(config.src_ip));
		memcpy_s(addrs.des_mac_addr, sizeof(addrs.des_mac_addr), config.gtw_mac, sizeof(config.gtw_mac));
		memcpy_s(addrs.des_ip_addr, sizeof(addrs.des_ip_addr), config.dns_server, sizeof(config.dns_server));

		thrd_t thr;
		// 将thr_fun函数放在一个新的线程中执行
		if (thrd_create(&thr, start_capture, &cnd) == thrd_error)
		{
			printf("ERORR; thrd_create() call failed\n");
			printf("按任意键退出...\n");
			getch(); // 等待用户输入
			exit(EXIT_FAILURE);
		}

		// thrd_sleep(&(struct timespec){.tv_sec = 1}, NULL); // 休眠1秒来等待抓包启动

		// 改为通过加锁和解锁来防止主线程先于子线程执行。
		// 以下三个步骤都有返回值来判断操作是否成功。
		// 方便起见就不处理错误了。
		mtx_lock(&mtx);		  // 加锁，阻塞主线程
		cnd_wait(&cnd, &mtx); // 等待子线程通知抓包启动。
		mtx_unlock(&mtx);	  // 解锁

		uint8_t res = send_dns_packet(pkt_id, config.s_port, &addrs, config.domain_name);

		if (res == SEND_DNS_PACKET_FAILED)
		{
			printf("Send DNS packet failed!\n");
		}

		printf("\n");

		// 此时子线程正在抓取并处理 DNS 服务器回复的数据包。

		int resa = thrd_join(thr, NULL); // 等待子线程结束
		if (resa == thrd_error)
		{
			printf("ERORR; thrd_join() call failed\n");
			printf("按任意键退出...\n");
			getch(); // 等待用户输入
			exit(EXIT_FAILURE);
		}

		return 0;
	}

	// 如果用户没有输入域名，则进入交互模式。
	// 定义一个长度为 100 的字符串接收控制台输入作为域名。
	char domain[100];

	interactive_mode = true;

	thrd_t thr;
	// 将thr_fun函数放在一个新的线程中执行
	if (thrd_create(&thr, start_capture, &cnd) == thrd_error)
	{
		printf("ERORR; thrd_create() call failed\n");
		printf("按任意键退出...\n");
		getch(); // 等待用户输入
		exit(EXIT_FAILURE);
	}

	mtx_lock(&mtx);		  // 加锁，阻塞主线程
	cnd_wait(&cnd, &mtx); // 等待子线程通知抓包已启动。
	mtx_unlock(&mtx);	  // 解锁

	printf("> ");

	while (1)
	{
		// 读取控制台输入，如果输入为空，则继续等待用户输入
		if (fgets(domain, 100, stdin) == NULL || domain[0] == '\n')
		{
			continue;
		}
		domain[strcspn(domain, "\n")] = 0; // 删除末尾的换行符

		// 如果用户输入 exit，则退出程序
		if (strcmp(domain, "exit") == 0)
		{
			exit(0);
		}

		// 如果用户输入 help，则打印帮助信息
		if (strcmp(domain, "help") == 0)
		{
			print_help();
			continue;
		}

		// 发送 DNS 查询数据包

		uint16_t pkt_id = PACKET_COUNTER;
		PACKET_COUNTER++;

		Addresses addrs;

		errno_t result;

		// 处理 memcpy_s 可能出现的错误。
		result = memcpy_s(addrs.source_mac_addr, sizeof(addrs.source_mac_addr), config.src_mac, sizeof(config.src_mac));
		if (result != 0)
		{
			fprintf(stderr, "memcpy_s failed: %d\n", result);
			printf("按任意键退出...\n");
			getch(); // 等待用户输入
			exit(1);
		}

		result = memcpy_s(addrs.source_ip_addr, sizeof(addrs.source_ip_addr), config.src_ip, sizeof(config.src_ip));
		if (result != 0)
		{
			fprintf(stderr, "memcpy_s failed: %d\n", result);
			printf("按任意键退出...\n");
			getch(); // 等待用户输入
			exit(1);
		}

		result = memcpy_s(addrs.des_mac_addr, sizeof(addrs.des_mac_addr), config.gtw_mac, sizeof(config.gtw_mac));
		if (result != 0)
		{
			fprintf(stderr, "memcpy_s failed: %d\n", result);
			printf("按任意键退出...\n");
			getch(); // 等待用户输入
			exit(1);
		}

		result = memcpy_s(addrs.des_ip_addr, sizeof(addrs.des_ip_addr), config.dns_server, sizeof(config.dns_server));
		if (result != 0)
		{
			fprintf(stderr, "memcpy_s failed: %d\n", result);
			exit(1);
		}

		uint8_t res = send_dns_packet(pkt_id, config.s_port, &addrs, domain);
		if (res == SEND_DNS_PACKET_FAILED)
		{
			printf("Send DNS packet failed!\n");
		}

		printf("\n");

		// 此时子线程正在抓取并处理 DNS 服务器回复的数据包。
	}

	return 0;
}
